package org.maplanning;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.TreeSet;

import org.maplanning.GridMap.StrategyI;

/**
 * Implements A* method for finding the best route. 
 * @author Jaak Simm
 */
public class AStarStrat implements StrategyI {
	/**
	 * A-star search algorithm, based on:
	 * http://en.wikipedia.org/wiki/A*_search_algorithm#Pseudocode
	 */
	@Override
	public char act(GridMap grid) {
		//TreeSet<Dist> x;
		
		HashSet<Node> closedSet = new HashSet<Node>(); // The set of nodes already evaluated. 
		HashSet<Node> openSet   = new HashSet<Node>(); // The set of tentative nodes to be evaluated, initially containing the start node
		HashMap<Node,Node> cameFrom = new HashMap<Node, Node>(); // The map of navigated nodes
		
		Matrix gScore = new Matrix(grid.height, grid.width); // Cost from start along best known path.
		Matrix hScore = new Matrix(grid.height, grid.width);
		Matrix fScore = new Matrix(grid.height, grid.width); // Estimated total cost from start to goal through y.
		
		// initializing A* search:
		openSet.add(new Node(grid.agent));
		gScore.set(grid.agent, 0);
		hScore.set(grid.agent, heuristic(grid.goalStates, grid.agent));
		fScore.set(grid.agent, gScore.get(grid.agent) + hScore.get(grid.agent));
		
		while ( !openSet.isEmpty() ) {
			// look up the node in openSet with the best score:
			Node current=null;
			double fScoreBest=Double.POSITIVE_INFINITY;
			for (Node n : openSet) {
				double v = fScore.get(n.x);
				if (v < fScoreBest) { // improving the best score:
					current = n;
					fScoreBest = v;
				}
			}
			
			// returning the action if reached goal:
			if (grid.get(current.x)==GridMap.GOAL) {
				ArrayList<int[]> path = reconstructPath(cameFrom, current);
				if (path.size()<=1) {
					return 0; // no action, already at goal
				}
				int[] x = path.get(path.size()-2).clone();
				int[] x2 = path.get(path.size()-1);
				x[0] -= x2[0];
				x[1] -= x2[1];
				// getting action would result in the first motion:
				return grid.getAction(x);
			}
			openSet.remove(current);
			closedSet.add(current);
			
			for (int[] neighbor : grid.getNeighbors(current.x)) {
				Node neighborNode = new Node(neighbor);
				if (closedSet.contains(neighborNode)) {
					continue;
				}
				double tentativeGScore = gScore.get(current.x) + 1;
				
				boolean tentativeIsBetter;
				// neighbor not in openSet:
				if (!openSet.contains(neighbor)) {
					openSet.add( neighborNode );
					hScore.set(neighbor, heuristic(grid.goalStates, neighbor));
					tentativeIsBetter = true;
				} else
				if (tentativeGScore < gScore.get(neighbor)) {
					tentativeIsBetter = true;
				} else {
					tentativeIsBetter = false;
				}
				
				// a new or better route to neighbor:
				if (tentativeIsBetter) {
					cameFrom.put(neighborNode, current);
					gScore.set(neighbor, tentativeGScore);
					fScore.set(neighbor, tentativeGScore+hScore.get(neighbor));
				}
			}
			
		}
		return 0;
	}
	
	/**
	 * @param cameFrom
	 * @param current
	 * @return path from current node to the beginning node
	 */
	static ArrayList<int[]> reconstructPath(HashMap<Node,Node> cameFrom, Node current) {
		ArrayList<int[]> path = new ArrayList<int[]>();
		
		path.add(current.x);
		// looping until reaching a node that is not in cameFrom:
		while (cameFrom.containsKey(current)) {
			current = cameFrom.get(current);
			path.add(current.x);
		}
		
		return path;
	}
	
	/**
	 * L1 heuristic from state to the closest goal
	 * @param goals
	 * @param state
	 * @return
	 */
	public static double heuristic(ArrayList<int[]> goals, int[] state) {
		double h = Double.POSITIVE_INFINITY;
		for (int[] goal : goals) {
			//double d = Math.sqrt(Math.pow(state[0]-goal[0], 2) + Math.pow(state[1]-goal[1], 2));
			double d = GridMap.distanceL1(state, goal);
			if (d < h) {
				h = d;
			}
		}
		return h;
	}
	
	/** Node, just stores int array, has proper equals() and hashCode() */
	class Node {
		int[] x;
		
		public Node(int[] x) { this.x = x; }
		
		@Override
		public boolean equals(Object obj) {
			if (!(obj instanceof Node)) {
				return false;
			}
			Node o = (Node)obj;
			return Arrays.equals(x, o.x);
		}
		@Override
		public int hashCode() {
			return Arrays.hashCode(this.x);
		}
		@Override
		public String toString() {
			return Arrays.toString(x);
		}
	}
	
	/**
	 * objects for storying distance to goal states
	 */
	class Dist implements Comparable<Dist> {
		int[] state;
		double distance;
		public Dist() {}
		
		@Override
		public int compareTo(Dist o) {
			if (this.distance!=o.distance) {
				return (int)Math.signum(this.distance - o.distance);
			}
			// for making compareTo consistent with equals:
			if (this.state[0]!=o.state[0]) {
				return this.state[0] - o.state[0];
			}
			return this.state[1] - o.state[1];
		}
		@Override
		public boolean equals(Object obj) {
			if (!(obj instanceof Dist)) {
				return false;
			}
			Dist o = (Dist)obj;
			return o.distance==distance && 
					o.state[0]==state[0] && 
					o.state[1]==state[1];
		}
	}
}
